# -*- coding=utf-8 -*-
"""
@author: xingwg
@license: (C) Copyright 2020-2025.
@contact: xingweiguo@chinasvt.com
@project: boya-reid
@file: naic.py
@time: 2020/9/12 19:40
@desc:
"""
import os
import glog
from collections import defaultdict
from .base_dataset import BaseImageDataset


class NAIC(BaseImageDataset):
    """
    2020 鹏城AI+行人重识别数据集
    """
    def __init__(self, data_dir="./data", verbose=True):
        super(NAIC, self).__init__()
        self.data_dir = data_dir
        # self.data_dir_train = os.path.join(self.data_dir, "train")
        self.data_dir_test = os.path.join(self.data_dir, "image_A")

        self.train_dataset = self._get_train_dataset(self.data_dir)
        self.query_dataset = self._get_test_dataset(self.data_dir_test, query=True)
        self.gallery_dataset = self._get_test_dataset(self.data_dir_test, query=False)

        if verbose:
            glog.info("=> NAIC data loaded.")
            self.print_dataset_statistics(self.train_dataset, self.query_dataset, self.gallery_dataset)

        self.num_train_person_ids, self.num_train_images, self.num_train_cam_ids = \
            self.get_dataset_info(self.train_dataset)

    @staticmethod
    def _get_train_dataset(data_dir):
        file_path = os.path.join(data_dir, "train_list.txt")
        if not os.path.exists(file_path):
            raise IOError("{} does not exist".format(file_path))

        dataset = []
        cam_id = 1
        count_images = defaultdict(list)
        with open(file_path, "r") as f:
            while True:
                lines = f.readline().strip()
                if not lines:
                    break
                image_name, image_label = [i for i in lines.split()]
                count_images[image_label].append(image_name)

        val_images = {}
        person_ids = set()
        for person_id, image_names in count_images.items():
            if len(image_names) < 2:
                continue
            val_images[person_id] = count_images[person_id]
            person_ids.add(person_id)
        pid2label = {pid: label for label, pid in enumerate(person_ids)}
        for pid, image_names in val_images.items():
            label = pid2label[pid]
            for image_name in image_names:
                dataset.append((os.path.join(data_dir, image_name), label, cam_id))

        return dataset

    @staticmethod
    def _get_test_dataset(data_dir, query=True):
        subfix = "query"
        if not query:
            subfix = "gallery"

        file_path = os.path.join(data_dir, "{}.txt".format(subfix))
        dataset = []
        with open(file_path, "r") as f:
            while True:
                line = f.readline().strip()
                if not line:
                    break
                image_name = line
                dataset.append((os.path.join(data_dir, subfix, image_name), 1, 1))

        return dataset
